{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Homework and bake-off: pragmatic color descriptions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "__author__ = \"Christopher Potts\"\n",
    "__version__ = \"CS224u, Stanford, Fall 2020\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contents\n",
    "\n",
    "1. [Overview](#Overview)\n",
    "1. [Set-up](#Set-up)\n",
    "1. [All two-word examples as a dev corpus](#All-two-word-examples-as-a-dev-corpus)\n",
    "1. [Dev dataset](#Dev-dataset)\n",
    "1. [Random train–test split for development](#Random-train–test-split-for-development)\n",
    "1. [Question 1: Improve the tokenizer [1 point]](#Question-1:-Improve-the-tokenizer-[1-point])\n",
    "1. [Use the tokenizer](#Use-the-tokenizer)\n",
    "1. [Question 2: Improve the color representations [1 point]](#Question-2:-Improve-the-color-representations-[1-point])\n",
    "1. [Use the color representer](#Use-the-color-representer)\n",
    "1. [Initial model](#Initial-model)\n",
    "1. [Question 3: GloVe embeddings [1 point]](#Question-3:-GloVe-embeddings-[1-point])\n",
    "1. [Try the GloVe representations](#Try-the-GloVe-representations)\n",
    "1. [Question 4: Color context [3 points]](#Question-4:-Color-context-[3-points])\n",
    "1. [Your original system [3 points]](#Your-original-system-[3-points])\n",
    "1. [Bakeoff [1 point]](#Bakeoff-[1-point])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "This homework and associated bake-off are oriented toward building an effective system for generating color descriptions that are pragmatic in the sense that they would help a reader/listener figure out which color was being referred to in a shared context consisting of a target color (whose identity is known only to the describer/speaker) and a set of distractors.\n",
    "\n",
    "The notebook [colors_overview.ipynb](colors_overview.ipynb) should be studied before work on this homework begins. That notebook provides backgroud on the task, the dataset, and the modeling code that you will be using and adapting.\n",
    "\n",
    "The homework questions are more open-ended than previous ones have been. Rather than asking you to implement pre-defined functionality, they ask you to try to improve baseline components of the full system in ways that you find to be effective. As usual, this culminates in a prompt asking you to develop a novel system for entry into the bake-off. In this case, though, the work you do for the homework will likely be directly incorporated into that system (not required, but an efficient way to work at the very least)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set-up"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See [colors_overview.ipynb](colors_overview.ipynb) for set-up in instructions and other background details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from colors import ColorsCorpusReader\n",
    "from nltk.translate.bleu_score import corpus_bleu\n",
    "import numpy as np\n",
    "import os\n",
    "from sklearn.model_selection import train_test_split\n",
    "from torch_color_describer import (\n",
    "    ContextualColorDescriber, create_example_dataset)\n",
    "import utils\n",
    "from utils import START_SYMBOL, END_SYMBOL, UNK_SYMBOL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "utils.fix_random_seeds()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "COLORS_SRC_FILENAME = os.path.join(\n",
    "    \"data\", \"colors\", \"filteredCorpus.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## All two-word examples as a dev corpus\n",
    "\n",
    "So that you don't have to sit through excessively long training runs during development, I suggest working with the two-word-only subset of the corpus until you enter into the late stages of system testing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev_corpus = ColorsCorpusReader(\n",
    "    COLORS_SRC_FILENAME,\n",
    "    word_count=2,\n",
    "    normalize_colors=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev_examples = list(dev_corpus.read())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This subset has about one-third the examples of the full corpus:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(dev_examples)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We __should__ worry that it's not a fully representative sample. Most of the descriptions in the full corpus are shorter, and a large proportion are longer. So this dataset is mainly for debugging, development, and general hill-climbing. All findings should be validated on the full dataset at some point."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dev dataset\n",
    "\n",
    "The first step is to extract the raw color and raw texts from the corpus:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev_rawcols, dev_texts = zip(*[[ex.colors, ex.contents] for ex in dev_examples])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The raw color representations are suitable inputs to a model, but the texts are just strings, so they can't really be processed as-is. Question 1 asks you to do some tokenizing!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Random train–test split for development\n",
    "\n",
    "For the sake of development runs, we create a random train–test split:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev_rawcols_train, dev_rawcols_test, dev_texts_train, dev_texts_test = \\\n",
    "    train_test_split(dev_rawcols, dev_texts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Question 1: Improve the tokenizer [1 point]\n",
    "\n",
    "This is the first required question – the first required modification to the default pipeline.\n",
    "\n",
    "The function `tokenize_example` simply splits its string on whitespace and adds the required start and end symbols:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize_example(s):\n",
    "\n",
    "    # Improve me!\n",
    "\n",
    "    return [START_SYMBOL] + s.split() + [END_SYMBOL]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenize_example(dev_texts_train[376])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "__Your task__: Modify `tokenize_example` so that it does something more sophisticated with the input text. \n",
    "\n",
    "__Notes__:\n",
    "\n",
    "* There are useful ideas for this in [Monroe et al. 2017](https://transacl.org/ojs/index.php/tacl/article/view/1142)\n",
    "* There is no requirement that you do word-level tokenization. Sub-word and multi-word are options.\n",
    "* This question can interact with the size of your vocabulary (see just below), and in turn with decisions about how to use `UNK_SYMBOL`.\n",
    "\n",
    "__Important__: don't forget to add the start and end symbols, else the resulting models will definitely be terrible! The following test will check that your tokenizer has this property:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_tokenize_example(func):\n",
    "    s = \"A test string\"\n",
    "    result = func(s)\n",
    "    assert all(isinstance(tok, str) for tok in result), \\\n",
    "        \"The tokenizer must return a list of strings.\"\n",
    "    assert result[0] == START_SYMBOL, \\\n",
    "        \"The tokenizer must add START_SYMBOL as the first token.\"\n",
    "    assert result[-1] == END_SYMBOL, \\\n",
    "        \"The tokenizer must add END_SYMBOL as the final token.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_tokenize_example(tokenize_example)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use the tokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the tokenizer is working, run the following cell to tokenize your inputs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev_seqs_train = [tokenize_example(s) for s in dev_texts_train]\n",
    "\n",
    "dev_seqs_test = [tokenize_example(s) for s in dev_texts_test]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use only the train set to derive a vocabulary for the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev_vocab = sorted({w for toks in dev_seqs_train for w in toks})\n",
    "\n",
    "dev_vocab += [UNK_SYMBOL]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's important that the `UNK_SYMBOL` is included somewhere in this list. In test examples, words not seen in training will be mapped to `UNK_SYMBOL`. \n",
    "\n",
    "Conceptual note: If you model's vocab is the same as your train vocab, then `UNK_SYMBOL` will never be encountered during training, so it will be a random vector at test time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(dev_vocab)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Question 2: Improve the color representations [1 point]\n",
    "\n",
    "This is the second required pipeline improvement for the assignment. \n",
    "\n",
    "The following functions do nothing at all to the raw input colors we get from the corpus. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def represent_color_context(colors):\n",
    "\n",
    "    # Improve me!\n",
    "\n",
    "    return [represent_color(color) for color in colors]\n",
    "\n",
    "\n",
    "def represent_color(color):\n",
    "\n",
    "    # Improve me!\n",
    "\n",
    "    return color"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "represent_color_context(dev_rawcols_train[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "__Your task__: Modify `represent_color_context` and/or `represent_color` to represent colors in a new way.\n",
    "    \n",
    "__Notes__:\n",
    "\n",
    "* You are not required to keep `represent_color`. This might be unnatural if you want to perform an operation on each color trio all at once.\n",
    "* For that matter, if you want to process all of the color contexts in the entire data set all at once, that is fine too, as long as you can also perform the operation at test time with an unknown number of examples being tested.\n",
    "\n",
    "* The Fourier-transform method of [Monroe et al. 2017](https://transacl.org/ojs/index.php/tacl/article/view/1142) is a proven choice for our task. __It is not required that you implement this.__ However, if you decide to, you might find that the overly terse presentation in the paper is an obstacle. They key thing to see is that the notation $\\hat{f}_{jkl}$ is meant to specify a full coordinate system. Thus, you might do something like\n",
    "\n",
    "  ```\n",
    "from itertools import product\n",
    "for j, k, l in product((0, 1, 2), repeat=3):    \n",
    "    f_jkl = ...\n",
    "```\n",
    "and collect these `f_jkl` values in a list of 27 values. Additionally, in Python, [`2j` produces a value with `real` and `imag` attributes](https://docs.python.org/3.7/library/cmath.html). Each element `f_jkl` should have these components. If you concatenate the `real` and `imag` parts of all the `f_jkl`, you will have a 54-dimensional representation, as in the paper. Remember to start with an HSV representation, and with $h$ in $[0, 360]$, $s$ in $[0, 200]$, and $v$ in $[0, 200]$ (or else do the scaling differently). Note that the values in our corpus are in HLS format, [which are easily converted to HSV](https://en.wikipedia.org/wiki/HSL_and_HSV#HSV_to_HSL).\n",
    "\n",
    "The following test seeks to ensure only that the output of your `represent_color_context` will be compatible with the models we are creating:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_represent_color_context(func):\n",
    "    \"\"\"`func` should be `represent_color_context`\"\"\"\n",
    "    example = [\n",
    "        [0.786, 0.58, 0.87],\n",
    "        [0.689, 0.44, 0.92],\n",
    "        [0.628, 0.32, 0.81]]\n",
    "    result = func(example)\n",
    "    assert len(result) == len(example), \\\n",
    "        (\"Color context representations need to represent each color \"\n",
    "         \"separately. (We assume the final color is the target.)\")\n",
    "    for i, color in enumerate(result):\n",
    "        assert all(isinstance(x, float) for x in color), \\\n",
    "            (\"All color representations should be lists of floats. \"\n",
    "             \"Color {} is {}\".format(i, color))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_represent_color_context(represent_color_context)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use the color representer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following cell just runs your `represent_color_context` on the train and test sets:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev_cols_train = [represent_color_context(colors) for colors in dev_rawcols_train]\n",
    "\n",
    "dev_cols_test = [represent_color_context(colors) for colors in dev_rawcols_test]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At this point, our preprocessing steps are complete, and we can fit a first model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initial model\n",
    "\n",
    "The first model is configured right now to be a small model run for just a few iterations. It should be enough to get traction, but it's unlikely to be a great model. You are free to modify this configuration if you wish; it is here just for demonstration and testing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev_mod = ContextualColorDescriber(\n",
    "    dev_vocab,\n",
    "    early_stopping=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if 'IS_GRADESCOPE_ENV' not in os.environ:\n",
    "    %time _ = dev_mod.fit(dev_cols_train, dev_seqs_train)\n",
    "else:\n",
    "    dev_mod.fit(dev_cols_train, dev_seqs_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The canonical bake-off evaluation function is `evaluate`. Our primary metric is `listener_accuracy`; the BLEU score is included as a check to ensure that your system is speaking English!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev_mod.evaluate(dev_cols_test, dev_seqs_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also see the model's predicted sequences given color context inputs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev_mod.predict(dev_cols_test[:1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev_seqs_test[:1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Question 3: GloVe embeddings [1 point]\n",
    "\n",
    "The above model uses a random initial embedding, as configured by the decoder used by `ContextualColorDescriber`. This homework question asks you to consider using GloVe inputs. \n",
    "\n",
    "__Your task__: Complete `create_glove_embedding` so that it creates a GloVe embedding based on your model vocabulary. This isn't mean to be analytically challenging, but rather just to create a basis for you to try out other kinds of rich initialization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "GLOVE_HOME = os.path.join('data', 'glove.6B')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_glove_embedding(vocab, glove_base_filename='glove.6B.50d.txt'):\n",
    "    pass\n",
    "    # Use `utils.glove2dict` to read in the GloVe file:\n",
    "\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "    # Use `utils.create_pretrained_embedding` to create the embedding.\n",
    "    # This function will, by default, ensure that START_TOKEN,\n",
    "    # END_TOKEN, and UNK_TOKEN are included in the embedding.\n",
    "\n",
    "    ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "    # Be sure to return the embedding you create as well as the\n",
    "    # vocabulary returned by `utils.create_pretrained_embedding`,\n",
    "    # which is likely to have been modified from the input `vocab`.\n",
    "\n",
    "    ##### YOUR CODE HERE\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_create_glove_embedding(func):\n",
    "    vocab = ['NLU', 'is', 'the', 'future', '.', '$UNK', '<s>', '</s>']\n",
    "    glove_embedding, glove_vocab = func(vocab, 'glove.6B.50d.txt')\n",
    "    assert isinstance(glove_embedding, np.ndarray), \\\n",
    "        \"Expected embedding type {}; got {}\".format(\n",
    "        glove_embedding.__class__.__name__, glove_embedding.__class__.__name__)\n",
    "    assert glove_embedding.shape == (8, 50), \\\n",
    "        \"Expected embedding shape (8, 50); got {}\".format(glove_embedding.shape)\n",
    "    assert glove_vocab == vocab, \\\n",
    "        \"Expected vocab {}; got {}\".format(vocab, glove_vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_create_glove_embedding(create_glove_embedding)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Try the GloVe representations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The extent to which GloVe is useful will depend heavily on how aligned your tokenization scheme is with the GloVe vocabulary. For example, if you did character-level tokenization, then the GloVe embedding space is not well-aligned with your tokenizer and using GloVe should have little no positive effect."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see if GloVe helped for our development data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev_glove_embedding, dev_glove_vocab = create_glove_embedding(dev_vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(dev_vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(dev_glove_vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev_mod_glove = ContextualColorDescriber(\n",
    "    dev_glove_vocab,\n",
    "    embedding=dev_glove_embedding,\n",
    "    early_stopping=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = dev_mod_glove.fit(dev_cols_train, dev_seqs_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev_mod_glove.listener_accuracy(dev_cols_test, dev_seqs_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You probably saw a small boost, assuming your tokeization scheme leads to good overlap with the GloVe vocabulary. The input representations are larger than in our previous model (at least as I configured things), so we would need to do more runs with higher `max_iter` values to see whether this is worthwhile overall."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Question 4: Color context [3 points]\n",
    "\n",
    "The final required homework question is the most challenging, but it should set you up to think in much more flexible ways about the underlying model we're using.\n",
    "\n",
    "The question asks you to modify various model components in `torch_color_describer.py`. The section called [Modifying the core model](colors_overview.ipynb#Modifying-the-core-model) from the core unit notebook provides a number of examples illustrating the basic techniques, so you might review that material if you get stuck here.\n",
    "\n",
    "__Your task__: Building on ideas from [Monroe et al. 2017](https://transacl.org/ojs/index.php/tacl/article/view/1142), you will redesign the model so that the target color (the final one in the context) is appended to each input token that gets processed by the decoder. The question asks you to subclass the `Decoder` and `EncoderDecoder` from `torch_color_describer.py` so that you can build models that do this."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "__Step 1__: Modify the `Decoder` so that the input vector to the model at each timestep is not just a token representation `x` but the concatenation of `x` with the representation of the target color.\n",
    "\n",
    "__Notes__:\n",
    "\n",
    "* You might notice at this point that the original `Decoder.forward` method has an optional keyword argument `target_colors` that is passed to `Decoder.get_embeddings`. Because this is already in place, all you have to do is modify the `get_embeddings` method to use this argument.\n",
    "\n",
    "* The change affects the configuration of `self.rnn`, so you need to subclass the `__init__` method as well, so that its `input_size` argument accomodates the embedding as well as the color representations.\n",
    "\n",
    "* You can do the relevant operations efficiently in pure PyTorch using `repeat_interleave` and `cat`, but the important thing is to get a working implementation – you can always optimize the code later if the ideas prove useful to you. \n",
    "\n",
    "Here's skeleton code for you to flesh out:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_color_describer import Decoder\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "class ColorContextDecoder(Decoder):\n",
    "    def __init__(self, color_dim, *args, **kwargs):\n",
    "        self.color_dim = color_dim\n",
    "        super().__init__(*args, **kwargs)\n",
    "\n",
    "        # Fix the `self.rnn` attribute:\n",
    "        ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "    def get_embeddings(self, word_seqs, target_colors=None):\n",
    "        \"\"\"\n",
    "        You can assume that `target_colors` is a tensor of shape\n",
    "        (m, n), where m is the length of the batch (same as\n",
    "        `word_seqs.shape[0]`) and n is the dimensionality of the\n",
    "        color representations the model is using. The goal is\n",
    "        to attached each color vector i to each of the tokens in\n",
    "        the ith sequence of (the embedded version of) `word_seqs`.\n",
    "\n",
    "        \"\"\"\n",
    "        ##### YOUR CODE HERE\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Step 1 is the most demanding of the steps in terms of tensor wrangling. It's important to have a clear idea of what you are trying to achieve and to unit test `get_embeddings` so that you can check that it has realized your vision. The following test should help with that:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_get_embeddings(decoder_class):\n",
    "    \"\"\"\n",
    "    It's assumed that the input to this will be `ColorContextDecoder`.\n",
    "    You pass in the class, and the function initalizes it with the test\n",
    "    parameters.\n",
    "    \"\"\"\n",
    "    dec = decoder_class(\n",
    "        color_dim=3,   # For these, we mainly want *different*\n",
    "        vocab_size=10, # dimensions so that we reliably get\n",
    "        embed_dim=4,   # dimensionality errors if something\n",
    "        hidden_dim=5)  # isn't working.\n",
    "\n",
    "    # This step just changes the embedding to one with values\n",
    "    # that are easy to inspect and definitely will not change\n",
    "    # between runs:\n",
    "    dec.embedding = nn.Embedding.from_pretrained(\n",
    "        torch.FloatTensor([\n",
    "            [10, 11, 12, 13],\n",
    "            [14, 15, 16, 17],\n",
    "            [18, 19, 20, 21]]))\n",
    "\n",
    "    # These are the incoming sequences -- lists of indices\n",
    "    # into the rows of `dec.embedding`:\n",
    "    word_seqs = torch.tensor([\n",
    "        [0,1,2],\n",
    "        [2,0,1]])\n",
    "\n",
    "    # Target colors as small floats that will be easy to track:\n",
    "    target_colors = torch.tensor([\n",
    "        [0.1, 0.2, 0.3],\n",
    "        [0.7, 0.8, 0.9]])\n",
    "\n",
    "    # The desired return value: one list of tensors for each of\n",
    "    # the two sequences in `word_seqs`. Each index is replaced\n",
    "    # with its vector from `dec.embedding` and has the\n",
    "    # corrresponding color from `target_colors` appended to it.\n",
    "    expected = torch.tensor([\n",
    "        [[10., 11., 12., 13.,  0.1,  0.2,  0.3],\n",
    "         [14., 15., 16., 17.,  0.1,  0.2,  0.3],\n",
    "         [18., 19., 20., 21.,  0.1,  0.2,  0.3]],\n",
    "\n",
    "        [[18., 19., 20., 21.,  0.7,  0.8,  0.9],\n",
    "         [10., 11., 12., 13.,  0.7,  0.8,  0.9],\n",
    "         [14., 15., 16., 17.,  0.7,  0.8,  0.9]]])\n",
    "\n",
    "    result = dec.get_embeddings(word_seqs, target_colors=target_colors)\n",
    "\n",
    "    assert expected.shape == result.shape, \\\n",
    "        \"Expected shape {}; got shape {}\".format(expected.shape, result.shape)\n",
    "\n",
    "    assert torch.all(expected.eq(result)), \\\n",
    "        (\"Your result has the desired shape but the values aren't correct. \"\n",
    "         \"Here's what your function creates; compare it with `expected` \"\n",
    "         \"from the test:\\n{}\".format(result))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_get_embeddings(ColorContextDecoder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "__Step 2__: Modify the `EncoderDecoder`. For this, you just need to make a small change to the `forward` method: extract the target colors from `color_seqs` and feed them to the decoder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_color_describer import EncoderDecoder\n",
    "\n",
    "class ColorizedEncoderDecoder(EncoderDecoder):\n",
    "\n",
    "    def forward(self,\n",
    "            color_seqs,\n",
    "            word_seqs,\n",
    "            seq_lengths=None,\n",
    "            hidden=None,\n",
    "            targets=None):\n",
    "        if hidden is None:\n",
    "            hidden = self.encoder(color_seqs)\n",
    "\n",
    "        # Extract the target colors from `color_seqs` and\n",
    "        # feed them to the decoder, which already has a\n",
    "        # `target_colors` keyword.\n",
    "\n",
    "        ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "        # Your decoder will return `output, hidden` pairs; the\n",
    "        # following will handle the two return situations that\n",
    "        # the code needs to consider -- training and prediction.\n",
    "        if self.training:\n",
    "            return output\n",
    "        else:\n",
    "            return output, hidden"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "__Step 3__: Finally, as in the examples in [Modifying the core model](colors_overview.ipynb#Modifying-the-core-model), you need to modify the `build_graph` method of `ContextualColorDescriber` so that it uses your new `ColorContextDecoder` and `ColorizedEncoderDecoder`. Here's starter code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_color_describer import Encoder\n",
    "\n",
    "class ColorizedInputDescriber(ContextualColorDescriber):\n",
    "\n",
    "    def build_graph(self):\n",
    "\n",
    "        # We didn't modify the encoder, so this is\n",
    "        # just copied over from the original:\n",
    "        encoder = Encoder(\n",
    "            color_dim=self.color_dim,\n",
    "            hidden_dim=self.hidden_dim)\n",
    "\n",
    "        # Use your `ColorContextDecoder`, making sure\n",
    "        # to pass in all the keyword arguments coming\n",
    "        # from `ColorizedInputDescriber`:\n",
    "\n",
    "        ##### YOUR CODE HERE\n",
    "\n",
    "\n",
    "\n",
    "        # Return a `ColorizedEncoderDecoder` that uses\n",
    "        # your encoder and decoder:\n",
    "\n",
    "        ##### YOUR CODE HERE\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's it! Since these modifications are pretty intricate, you might want to use [a toy dataset](colors_overview.ipynb#Toy-problems-for-development-work) to debug it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_full_system(describer_class):\n",
    "    toy_color_seqs, toy_word_seqs, toy_vocab = create_example_dataset(\n",
    "        group_size=50, vec_dim=2)\n",
    "\n",
    "    toy_color_seqs_train, toy_color_seqs_test, toy_word_seqs_train, toy_word_seqs_test = \\\n",
    "        train_test_split(toy_color_seqs, toy_word_seqs)\n",
    "\n",
    "    toy_mod = describer_class(toy_vocab)\n",
    "\n",
    "    _ = toy_mod.fit(toy_color_seqs_train, toy_word_seqs_train)\n",
    "\n",
    "    acc = toy_mod.listener_accuracy(toy_color_seqs_test, toy_word_seqs_test)\n",
    "\n",
    "    return acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_full_system(ColorizedInputDescriber)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If that worked, then you can now try this model on SCC problems!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Your original system [3 points]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are many options for your original system, which consists of the full pipeline – all preprocessing and modeling steps. You are free to use any model you like, as long as you subclass `ContextualColorDescriber` in a way that allows its `evaluate` method to behave in the expected way.\n",
    "\n",
    "So that we can evaluate models in a uniform way for the bake-off, we ask that you modify the function `my_original_system` below so that it accepts a trained instance of your model and does any preprocessing steps required by your model.\n",
    "\n",
    "If we seek to reproduce your results, we will rerun this entire notebook. Thus, it is fine if your `my_original_system` makes use of functions you wrote or modified above this cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def my_original_system(trained_model, color_seqs_test, texts_test):\n",
    "    \"\"\"\n",
    "    Feel free to modify this code to accommodate the needs of\n",
    "    your system. Just keep in mind that it will get raw corpus\n",
    "    examples as inputs for the bake-off.\n",
    "\n",
    "    \"\"\"\n",
    "    # `word_seqs_test` is a list of strings, so tokenize each of\n",
    "    # its elements:\n",
    "    tok_seqs = [tokenize_example(s) for s in texts_test]\n",
    "\n",
    "    col_seqs = [represent_color_context(colors)\n",
    "                for colors in color_seqs_test]\n",
    "\n",
    "\n",
    "    # Optionally include other preprocessing steps here. Note:\n",
    "    # DO NOT RETRAIN YOUR MODEL! It's a tempting step, but it's\n",
    "    # a mistake and will get you disqualified!\n",
    "\n",
    "\n",
    "    # The following core score calculations are required:\n",
    "    return trained_model.evaluate(col_seqs, tok_seqs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If `my_original_system` works on test sets you create from the corpus distribution, then it will work for the bake-off, so consider checking that. For example, this would check that `dev_mod` above passes muster:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_original_system(dev_mod, dev_rawcols_test, dev_texts_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the cell below, please provide a brief technical description of your original system, so that the teaching team can gain an understanding of what it does. This will help us to understand your code and analyze all the submissions to identify patterns and strategies. We also ask that you report the best **listener_accuracy** score your system got during development, just to help us understand how systems performed overall."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PLEASE MAKE SURE TO INCLUDE THE FOLLOWING BETWEEN THE START AND STOP COMMENTS:\n",
    "#   1) Textual description of your system.\n",
    "#   2) The code for your original system.\n",
    "#   3) The score achieved by your system in place of MY_NUMBER.\n",
    "#        With no other changes to that line.\n",
    "#        You should report your score as a decimal value <=1.0\n",
    "# PLEASE MAKE SURE NOT TO DELETE OR EDIT THE START AND STOP COMMENTS\n",
    "\n",
    "# NOTE: MODULES, CODE AND DATASETS REQUIRED FOR YOUR ORIGINAL SYSTEM \n",
    "# SHOULD BE ADDED BELOW THE 'IS_GRADESCOPE_ENV' CHECK CONDITION. DOING\n",
    "# SO ABOVE THE CHECK MAY CAUSE THE AUTOGRADER TO FAIL.\n",
    "\n",
    "# START COMMENT: Enter your system description in this cell.\n",
    "# My peak score was: MY_NUMBER\n",
    "if 'IS_GRADESCOPE_ENV' not in os.environ:\n",
    "    pass\n",
    "\n",
    "# STOP COMMENT: Please do not remove this comment."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bakeoff [1 point]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the bake-off, we will release a test set. You will evaluate your custom model from the previous question on these new datasets using your `my_original_system` function. Rules:\n",
    "\n",
    "1. Only one evaluation is permitted.\n",
    "1. No additional system tuning is permitted once the bake-off has started.\n",
    "\n",
    "Systems will be ranked primarily by `listener_accuracy`, but we will also consider their `corpus_bleu` scores as reported by `my_original_system`. However, the BLEU score is just a simple check that your system is speaking some version of English that corresponds in some meaningful way to the gold descriptions, so you should concentrate on `listener_accuracy`.\n",
    "\n",
    "The cells below this one constitute your bake-off entry.\n",
    "\n",
    "People who enter will receive the additional homework point, and people whose systems achieve the top score will receive an additional 0.5 points. We will test the top-performing systems ourselves, and only systems for which we can reproduce the reported results will win the extra 0.5 points.\n",
    "\n",
    "Late entries will be accepted, but they cannot earn the extra 0.5 points. Similarly, you cannot win the bake-off unless your homework is submitted on time.\n",
    "\n",
    "The announcement will include the details on where to submit your entry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Enter your bake-off assessment code in this cell.\n",
    "# Please do not remove this comment.\n",
    "if 'IS_GRADESCOPE_ENV' not in os.environ:\n",
    "    pass\n",
    "    # Please enter your code in the scope of the above conditional.\n",
    "    ##### YOUR CODE HERE\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# On an otherwise blank line in this cell, please enter\n",
    "# the return value of `my_original_system` from your\n",
    "# bake-off run. You can just paste it in.\n",
    "# Please do not remove this comment.\n",
    "if 'IS_GRADESCOPE_ENV' not in os.environ:\n",
    "    pass\n",
    "    # Please paste in your results in the scope of the above conditional.\n",
    "    ##### YOUR SCORE HERE\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
